import pygame
import numpy as np

# 1. 初始化参数
WIDTH, HEIGHT = 900, 700  # 窗口尺寸
CELL_SIZE = 10  # 每个元胞的像素大小
GRID_WIDTH = WIDTH // CELL_SIZE  # 网格列数
GRID_HEIGHT = HEIGHT // CELL_SIZE  # 网格行数
FPS = 10  # 每秒刷新帧数（控制动画速度）

# 颜色定义
BLACK = (0, 0, 0)  # 死亡元胞
WHITE = (255, 255, 255)  # 存活元胞
GRAY = (128, 128, 128)  # 网格线

# 2. 初始化Pygame和窗口
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame 元胞自动机（康威生命游戏）")
clock = pygame.time.Clock()

# 3. 初始化网格（随机生成存活/死亡状态，True=存活，False=死亡）
grid = np.random.choice([True, False], size=(GRID_WIDTH, GRID_HEIGHT), p=[0.2, 0.8])


# 4. 计算相邻存活元胞数量的核心函数
def count_neighbors(x, y):
    count = 0
    # 遍历3x3邻域（排除自身）
    for dx in (-1, 0, 1):
        for dy in (-1, 0, 1):
            if dx == 0 and dy == 0:
                continue  # 跳过当前元胞
            # 处理边界（循环边界：左右/上下相连）
            nx = (x + dx) % GRID_WIDTH
            ny = (y + dy) % GRID_HEIGHT
            if grid[nx, ny]:
                count += 1
    return count


# 5. 主循环
running = True
paused = False  # 暂停开关（默认不暂停）

while running:
    # 控制帧率
    clock.tick(FPS)

    # 事件处理（窗口关闭、鼠标点击、空格暂停）
    for event in pygame.event.get():
        # 关闭窗口
        if event.type == pygame.QUIT:
            running = False
        # 鼠标点击修改元胞状态（左键存活，右键死亡）
        if event.type == pygame.MOUSEBUTTONDOWN:
            x, y = pygame.mouse.get_pos()
            grid_x = x // CELL_SIZE
            grid_y = y // CELL_SIZE
            if event.button == 1:  # 左键
                grid[grid_x, grid_y] = True
            elif event.button == 3:  # 右键
                grid[grid_x, grid_y] = False
        # 空格键切换暂停/继续
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                paused = not paused

            # 如果点击空格暂停后点击c,即可清除屏幕所有元胞
            elif event.key == pygame.K_c:
                grid = np.full((GRID_WIDTH, GRID_HEIGHT), False)

    # 清空屏幕（黑色背景）
    screen.fill(BLACK)

    # 6. 更新和绘制网格（未暂停时才更新状态）
    new_grid = grid.copy()  # 新建网格存储下一状态（避免边更新边读取）
    for x in range(GRID_WIDTH):
        for y in range(GRID_HEIGHT):
            # 绘制元胞（存活=白色，死亡=黑色）
            rect = pygame.Rect(x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE - 1, CELL_SIZE - 1)
            pygame.draw.rect(screen, WHITE if grid[x, y] else BLACK, rect)

            # 未暂停时，根据规则更新状态
            if not paused:
                neighbors = count_neighbors(x, y)
                # 康威规则：
                # 1. 存活元胞：邻居<2或>3 → 死亡
                # 2. 死亡元胞：邻居=3 → 复活
                if grid[x, y]:
                    new_grid[x, y] = (neighbors == 2 or neighbors == 3)
                else:
                    new_grid[x, y] = (neighbors == 3)

    # 更新网格（仅在未暂停时替换）
    if not paused:
        grid = new_grid

    # 刷新屏幕显示
    pygame.display.flip()

# 退出程序
pygame.quit()